#include "Solver.hpp"

//smoothing function using Gauss-Seidel method
void Solver::smoothGS(Array &x, const Array &f, size_t iter)
{
    //assert(x.xsize() == f.xsize());

    size_t width = x.xsize();
    real h = x.hSize();
    //omega = 1.0 for Gauss-Seidel method
    real omg = 1.0;

    for ( size_t iter_ = 1 ; iter_<=iter ; ++iter_)
    {
        for (size_t i = 1; i < width-1; ++i)
        {
            for (size_t j = 1; j < width-1; ++j)
            {
	      x(i,j) = (1-omg)*x(i,j) + omg*h*h*0.25*( (x(i+1,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j) );
            }
        }
    }  
}

//smoothing function using red-black Gauss-Seidel
void Solver::smoothGSRB(Array &x, const Array &f, size_t iter)
{
    //assert(x.xSize() == f.xsize());

    size_t width = x.xsize();
    size_t j;
    real h = x.hSize();
    real omg = 1.0;


    for ( size_t iter_ = 1 ; iter_<=iter ; ++iter_)
    { 
        //updating "red" cells
            for (size_t i = 1; i < width-1; ++i)
            {
                (i%2 == 1)?j=1:j=2;
                for (; j < width-1; j+=2)
                {
		  x(i,j) = (1-omg)*x(i,j) + omg*h*h*0.25*( (x(i+1,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j) );
                }
            }
    //updating "black" cells
        for (size_t i = 1; i < width-1; ++i)
        {
            (i%2 == 1)?j = 2:j=1;
            for (; j < width-1; j+=2)
            {
	      x(i,j) = (1-omg)*x(i,j) + omg*h*h*0.25*( (x(i+1,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(h*h) -f(i,j) );
            }
        }
    }
}

void Solver::restrict(const Array &r, Array &f)
{
    //          assert ( r. xsize () == r. ysize () );
    //        assert ( r. xsize () == 2*f. xsize () -1 );
    //      assert ( r. ysize () == 2*f. ysize () -1 );
    //    assert ( real (2)* r. hsize () == f. hsize () );

    const size_t size ( f. xsize () );
    for ( size_t i =1; i<size -1; ++i ) 
    {
        for ( size_t j =1; j<size -1; ++j )
        {
            f(i,j) = I_.C  * r(i*2   , j*2  ) +
                I_.N  * r(i *2+1 , j*2  ) +
                I_.S  * r(i*2 -1 , j*2  ) +
                I_.W  * r(i*2   , j *2 -1) +
                I_.E  * r(i*2   , j *2+1) +
                I_.NW * r(i *2+1 , j *2 -1) +
                I_.NE * r(i *2+1 , j *2+1) +
                I_.SW * r(i*2 -1 , j *2 -1) +
                I_.SE * r(i*2 -1 , j *2+1);
        }
    }
}

//Ax-f =r (the residual is not elevated to the second power)
void Solver::calcResidual(const Array &x, const Array &f, Array &r)
{
    const size_t size ( f. xsize () );   
    for(size_t i=1;i<size-1;++i)
    {
        for(size_t j=1;j<size-1;++j)
        {
            r(i,j)=(x(i+1,j)-4.0*x(i,j)+x(i-1,j)+x(i,j+1)+x(i,j-1))/(x.hSize()*x.hSize())-f(i,j);
        }
    }
}


void Solver::setStencil()
{

    A_.N = A_.S = A_.W = A_.E = 1;
    A_.NW =A_.NE=A_.SW = A_.SE =0;
    A_.C = -4;	

    I_.N = I_.S = I_.W = I_.E = 2/16;
    I_.NW =I_.NE=I_.SW = I_.SE =1/16;
    I_.C = 4/16;

    J_.N = J_.S = J_.W = J_.E = 2/4;
    J_.NW =J_.NE=J_.SW = J_.SE =1/4;
    J_.C = 4/4;

}
void Solver::prolongateAndCorrect(const Array &e, Array &x)
{
    //          assert ( r. xsize () == r. ysize () );
    //        assert ( r. xsize () == 2*f. xsize () -1 );
    //      assert ( r. ysize () == 2*f. ysize () -1 );
    //    assert ( real (2)* r. hsize () == f. hsize () );

    // Prolongate Equation: e[level-1] = J * e[level-2];
    // Correction Equation: x[level-1] = x[level-1] + e[level-1]			

    const size_t size ( e. xsize () );
    for ( size_t i =1; i<size -1; ++i ) 
    {
        for ( size_t j =1; j<size -1; ++j )
        {
            x(2*i,2*j) += e(i,j);
            x(2*i+1,2*j) += (e(i,j)+e(i+1,j))*J_.E;
            x(2*i-1,2*j) += (e(i,j)+e(i-1,j))*J_.W;
            x(2*i,2*j+1) += (e(i,j)+e(i,j+1))*J_.N;
            x(2*i,2*j-1) += (e(i,j)+e(i,j-1))*J_.S;
            x(2*i+1,2*j+1) += (e(i,j)+e(i,j+1)+e(i+1,j)+e(i+1,j+1))*J_.NE;
            x(2*i-1,2*j+1) += (e(i,j)+e(i,j+1)+e(i-1,j)+e(i-1,j+1))*J_.NW;
            x(2*i+1,2*j-1) += (e(i,j)+e(i,j-1)+e(i+1,j)+e(i+1,j-1))*J_.SE;
            x(2*i-1,2*j-1) += (e(i,j)+e(i,j-1)+e(i-1,j)+e(i-1,j-1))*J_.SW;
        }
    }
}

void Solver::mgSolve(Array& x, const Array& f, size_t cycles, size_t pre_iter, size_t post_iter)
{
  size_t finest_l=levels_-1;
  for(size_t c=0;c<cycles;++c)
    {
      //Finest grid pre-smoothing
      smoothGSRB(x,f, pre_iter);
      //Finest grid residual calculation
      calcResidual(x,f,(*r_[finest_l]));
      //restriting the residual
      restrict((*r_[finest_l]), (*f_[finest_l-1]));
      //Initialize the next grid
      x_[finest_l-1]->fill(0.0);
        
      //intermediate grids
      for(size_t l=finest_l-1;l>0;--l)
        {
	  //pre-smoothing
	  smoothGSRB((*x_[l]),(*f_[l]), pre_iter);
	  //Finest grid residual calculation
	  calcResidual((*x_[l]),(*f_[l]),(*r_[l]));
	  //restriting the residual
	  restrict((*r_[l]),(*f_[l-1]));
	  //Initialize the next grid
	  x_[l-1]->fill(0.0);    
        }
      //Solving coarse grid 
      smoothGS((*x_[0]),(*f_[0]),1);
      //Intermediate levels
      for(size_t l=1;l<finest_l;++l)
        {
	  //Prolongate and correct 
	  prolongateAndCorrect((*x_[l-1]),(*x_[l]));
	  //Smooth
	  smoothGSRB((*x_[l]),*(f_[l]),post_iter);
        }
        
      //Prolongate and correct finest grid 
      prolongateAndCorrect((*x_[finest_l-1]),x);
      //Smooth finest grid
      smoothGSRB(x,f,post_iter);
      //Calculate residual
      calcResidual(x,f,(*r_[finest_l]));
      std::cout<<"Cycle "<<c+1<<" Convergence "<<r_[finest_l]->l2Norm()<<"\n";
        
    }
}
